#include <unistd.h> 
#include <chrono>
#include <iostream>
#include "NaluDebug.h"
#include "FileParse.hpp"

static const int MAX_NAL_SIZE = 1*1024*1024;

static int findStartcode3(unsigned char *buffer)
{
    return (buffer[0]==0 && buffer[1]==0 && buffer[2]==1);
}
static int findStartcode4(unsigned char *buffer)
{
    return (buffer[0]==0 && buffer[1]==0 && buffer[2]==0 && buffer[3]==1);
}

FileParse::FileParse(const char* filename)
: m_hH264(nullptr),
  m_hH265(nullptr),
  m_nType(FILE_UNK),
  max_nalu_size(0)
{
    auto start = std::chrono::high_resolution_clock::now();
    if(filename == nullptr){
        parse_error("filename is NUL\n");
        return;
    }
    if(access(filename, F_OK) == -1){
        parse_error("file not exist\n");
        return;
    }
    m_filename = std::string(filename);
    m_nType = judeVideoFile(m_filename.c_str());

    if (m_nType == FILE_H264){
        m_hH264 = h264_new();
    }else if (m_nType == FILE_H265){
        m_hH265 = h265_new();
    }else{
        parse_error("%s detect h264/h265 file error", m_filename.c_str());
        return ;
    }
    FILE* m_pFile = fopen(m_filename.c_str(), "r+b");
    if(m_pFile == nullptr){
        parse_error("fopen %s error \n", m_filename.c_str());
        return;
    }

    fseek(m_pFile, 0, SEEK_END); 
    size_t size = static_cast<size_t>(ftell(m_pFile));
    fseek(m_pFile, 0, SEEK_SET);
    m_vNalTypeVector.clear();
    parseFile(m_pFile, m_vNalTypeVector);
    mergeNalu();
    if (m_nType == FILE_H264){
        m_hH264->info->nalu_count = m_vNalTypeVector.size();
        m_hH264->info->file_size = size;
    }else{
        m_hH265->info->nalu_count = m_vNalTypeVector.size();
        m_hH265->info->file_size = size;
    }
    if (m_pFile != nullptr){
        fclose(m_pFile);
        m_pFile = nullptr;
    }
    auto end = std::chrono::high_resolution_clock::now();
    auto cost = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
    parse_info("[%s] parse cost: %ld ms frame cout: %lu\n",m_filename.c_str(), cost.count(), m_vNalTypeVector.size());
}

FileParse::Ptr FileParse::CreateShared(const char *filepath)
{
    // struct EnableMakeShared : public FileParse {
    //     EnableMakeShared(const char *filepath) : FileParse(filepath) {}
    // };
    // return std::make_shared<EnableMakeShared>(filepath);
    return Ptr(new FileParse(filepath));
}

FileParse::~FileParse()
{
    if (m_hH264 != nullptr){
        h264_free(m_hH264);
        m_hH264 = nullptr;
    }
    if (m_hH265 != nullptr){
        h265_free(m_hH265);
        m_hH265 = nullptr;
    }
    m_vNalTypeVector.clear();
    parse_info("delete FileParse end\n");
}

enum FileParse::FileType FileParse::GetFileType(const char* filename) 
{
    enum FileType type = FILE_UNK;
    const char* dot = strrchr(filename, '.');
    if (!dot || dot == filename) return FILE_UNK; // 没有扩展名
    
    // 检查文件扩展名
#if defined(_WIN32) || defined(_WIN64) || defined(WIN32) || defined(WIN64)
    if (!_stricmp(dot + 1, "h265") || !_stricmp(dot + 1, "265") || !_stricmp(dot + 1, "hevc")) {
        type = FILE_H265;
    } else if (!_stricmp(dot + 1, "h264") || !_stricmp(dot + 1, "264") || !_stricmp(dot + 1, "avc")) {
        type = FILE_H264;
    }
#else 
    if (!strcasecmp(dot + 1, "h265") || !strcasecmp(dot + 1, "265") || !strcasecmp(dot + 1, "hevc")) {
        type = FILE_H265;
    } else if (!strcasecmp(dot + 1, "h264") || !strcasecmp(dot + 1, "264") || !strcasecmp(dot + 1, "avc")) {
        type = FILE_H264;
    }
#endif 

    return type;
}

void FileParse::getVideoInfo(videoinfo_t* videoInfo)
{
    if(videoInfo == nullptr){
        std::cout << "videoInfo is nullptr" << std::endl;
        return;
    }
    if (m_nType == FILE_H265) {
        memcpy(videoInfo, m_hH265->info, sizeof(videoinfo_t));
        return;
    }
    else if(m_nType == FILE_H264) {
        memcpy(videoInfo, m_hH264->info, sizeof(videoinfo_t));
        return;
    }
    memset(videoInfo, 0, sizeof(videoinfo_t));
    videoInfo->type = FILE_UNK;
}

size_t FileParse::fileSize() const
{
    if (m_nType == FILE_H265) {
        return m_hH265->info->file_size;
    }
    else if(m_nType == FILE_H264) {
        return m_hH264->info->file_size;
    }
    return 0;
}

size_t FileParse::naluCount() const
{
    if (m_nType == FILE_H265) {
        return m_hH265->info->nalu_count;
    }
    else if(m_nType == FILE_H264) {
        return m_hH264->info->nalu_count;
    }
    return 0;
}

size_t FileParse::maxNaluSize() const
{
    return max_nalu_size;
}

int32_t FileParse::width() const
{
    if (m_nType == FILE_H265) {
        return m_hH265->info->width;
    }
    else if(m_nType == FILE_H264) {
        return m_hH264->info->width;
    }
    return 0;
}

int32_t FileParse::height() const
{
    if (m_nType == FILE_H265) {
        return m_hH265->info->height;
    }
    else if(m_nType == FILE_H264) {
        return m_hH264->info->height;
    }
    return 0;
}

int32_t FileParse::fileCodecId() const
{
    return m_nType;
}

std::string FileParse::fileCodecName() const
{
    std::string codec_name = "unknown";
    if (m_nType == FILE_H265) {
        codec_name = "H265";
    }
    else if(m_nType == FILE_H264) {
        codec_name = "H264";
    }
    return codec_name;
}

float FileParse::fps() const
{
    if (m_nType == FILE_H264) {
        return m_hH264->info->fps;
    }
    else if(m_nType == FILE_H265) {
        return m_hH265->info->fps;
    }
    return 0;
}

void FileParse::DebugSPS() const
{
    if (m_nType == FILE_H265) {
        h265_debug_sps(m_hH265->sps);
    }
    else if(m_nType == FILE_H264) {
        h264_debug_sps(m_hH264->sps);
    }
}

void FileParse::DebugPPS() const
{
    if (m_nType == FILE_H265) {
        h265_debug_pps(m_hH265->pps);
    }
    else if(m_nType == FILE_H264) {
        h264_debug_pps(m_hH264->pps);
    }
}

void FileParse::DebugVPS() const
{
    if (m_nType == FILE_H265) {
        h265_debug_vps(m_hH265->vps);
    }
}

void FileParse::mergeNalu()
{
    std::vector<NALU_t> vNal;
    size_t size = 0;
    int32_t first_nal_type = 0;
    for(size_t i = 0; i< m_vNalTypeVector.size(); i++){
        NALU_t n = m_vNalTypeVector[i];
        if(configFrame(&n)){
            if(size == 0){
                first_nal_type = n.nalType;
            }
            size += n.len;
            continue;
        }
        if(size > 0){
            if(dropAble(&n)){
                size += n.len;
                continue;
            }
            NALU_t new_nalu;
            memcpy(&new_nalu, &n, sizeof(NALU_t));
            new_nalu.offset = n.offset - size;
            size += n.len;
            new_nalu.nalType = first_nal_type;
            new_nalu.len = size;
            if(size > max_nalu_size){
                max_nalu_size = size;
            }
            new_nalu.num = vNal.size();
            vNal.push_back(new_nalu);
            size = 0;
            continue;
        }else{
            NALU_t new_nalu;
            memcpy(&new_nalu, &n, sizeof(NALU_t));
            new_nalu.num = vNal.size();
            if(size > max_nalu_size){
                max_nalu_size = size;
            }
            vNal.push_back(n);
            size = 0;
        }
    }
    m_vNalTypeVector.clear();
    m_vNalTypeVector.swap(vNal);
}

bool FileParse::getNaluInfo(size_t index, NALU_t *nalu)
{
    memset(nalu, 0 ,sizeof(NALU_t));
    if (index >= m_vNalTypeVector.size()){
        parse_error("index out of range\n");
        return false;
    }
    if(nalu == nullptr){
        parse_error("nalu nullptr error\n");
        return false;
    }
    memcpy(nalu, &m_vNalTypeVector[index], sizeof(NALU_t));
    return true;
}

size_t FileParse::getNaluData(size_t index, uint8_t *data, size_t srcSize)
{
    size_t dst_size = 0;
    if(data == nullptr){
        parse_error("src data nullptr error \n");
        return dst_size;
    }
    if (index >= m_vNalTypeVector.size()){
        parse_error("index(%lu) out of range(%lu) \n", index, m_vNalTypeVector.size());
        return dst_size;
    }
    FILE* m_pFile = fopen(m_filename.c_str(), "r+b");
    if(m_pFile == nullptr){
        parse_error("fopen %s error \n", m_filename.c_str());
        return dst_size;
    }
    NALU_t n = m_vNalTypeVector[index];
    if(n.len <= 0){
        fclose(m_pFile);
        m_pFile = nullptr;
        parse_error("%lu len 0 error \n", index);
        return dst_size;
    }
    if(srcSize < n.len){
        fclose(m_pFile);
        m_pFile = nullptr;
        parse_error("srcSize %lu < n.len %u error \n", srcSize, n.len);
        return dst_size;
    }
    fseek(m_pFile, n.offset, SEEK_SET);
    size_t ret = fread(data, 1, n.len, m_pFile);
    if(ret != n.len){
        parse_error("fread %lu != nalue->len %u \n",ret, n.len);
        fclose(m_pFile);
        m_pFile = nullptr;
        return dst_size;
    }
    dst_size = n.len;
    fclose(m_pFile);
    m_pFile = nullptr;
    return dst_size;
}

/**
解析NAL，返回两个开始字符之间间隔的字节数，即包含startcode的NALU的长度

note：一个视频文件中不同的NAL，startcode可能不一样。比如SPS为4字节，但SEI可能为3字节
todo:每次读一个字节，较慢，有无好的方法？
*/

int FileParse::parseFile(FILE* fp, std::vector<NALU_t> &nalus)
{
    if(fp == nullptr)return -1;
    fseek(fp, 0, SEEK_SET);
    constexpr size_t BUFFER_SIZE = 10 * 1024 * 1024;
    uint8_t *buffer = (uint8_t *)malloc(BUFFER_SIZE);
    if (buffer == nullptr) {
        fprintf(stderr, "Memory allocation failed\n");
        return -1;
    }
    nalus.clear();

    uint64_t file_position = 0;  // Current position in file
    uint64_t last_nalu_pos = 0;  // Start position of the last NALU found
    size_t bytes_read = 0;

    while ((bytes_read = fread(buffer, 1, BUFFER_SIZE, fp)) > 0) {
        for (size_t i = 0; i < bytes_read; i++) {
            if (i < bytes_read - 4 &&
                buffer[i] == 0x00 && buffer[i + 1] == 0x00 &&
                ((buffer[i + 2] == 0x01) || (buffer[i + 2] == 0x00 && buffer[i + 3] == 0x01))) {
                
                // Complete the previous NALU size calculation before recording new one
                if (nalus.size() > 0) {
                    nalus[nalus.size() - 1].len = file_position + i - last_nalu_pos;
                }

                // Record the current NALU header position
                last_nalu_pos = file_position + i;
                NALU_t new_nalu;
                new_nalu.offset = last_nalu_pos;
                int startcodeLen = (buffer[i + 2] == 0x01) ? 2 : 3;
                new_nalu.startcodeLen = startcodeLen + 1;
                new_nalu.type = m_nType;
                new_nalu.num = nalus.size();
                if(m_nType == 0){
                    new_nalu.nalType = buffer[i + new_nalu.startcodeLen] & 0x1f;
                    new_nalu.first_mb_in_slice = (buffer[i + new_nalu.startcodeLen + 1] & 0xff);
                }else{
                    new_nalu.nalType = (buffer[i + new_nalu.startcodeLen] >> 1) & 0x3f;
                    new_nalu.first_mb_in_slice = (buffer[i + new_nalu.startcodeLen + 2] & 0xff);
                }
                nalus.push_back(new_nalu);
                i += startcodeLen; // Move past the header
            }
        }
        file_position += bytes_read;
    }

    // Finalize the size of the last NALU
    if (nalus.size() > 0) {
        nalus[nalus.size() - 1].len = file_position - last_nalu_pos;
    }
    free(buffer);
    buffer = nullptr;
    for(size_t i = 0; i < nalus.size(); i++){
        parseNALU(fp, &nalus[i]);
    }
    fseek(fp, 0, SEEK_SET);
    return 0;
}

int FileParse::parseNALU(FILE* fp, NALU_t* nalu)
{
    if(fp == nullptr || nalu == nullptr){
        parse_error("fp == nullptr || nalu == nullptr\n");
        return -1;
    }
    if(nalu->len == 0 || (nalu->startcodeLen  != 3 && nalu->startcodeLen != 4)){
        parse_error("nalu->len == 0 || nalu->startcodeLen %d error\n", nalu->startcodeLen);
        return -1;
    }
    // parse_info("nalutype: %d %d %d \n", nalu->nalType, configFrame(nalu) ,keyframe(nalu));
    if(!configFrame(nalu) && !keyframe(nalu))return 0;
    size_t pos = ftell(fp);
    if(fseek(fp, nalu->offset, SEEK_SET) != 0){
        parse_error("fseek %u error \n", nalu->offset);
        return -1;
    }
    uint8_t *buffer = (uint8_t *)malloc(sizeof(uint8_t) * nalu->len);
    if(buffer == nullptr){
        parse_error("malloc error\n");
        return -1;
    }
    size_t size = fread(buffer, 1, nalu->len, fp);
    if(size != nalu->len){
        free(buffer);
        buffer = nullptr;
        parse_error("fread %lu != nalue->len %u \n",size, nalu->len);
        return -1;
    }
    if(m_nType == 0){
        m_hH264->sh->read_slice_type = 1;
        read_nal_unit(m_hH264, &buffer[nalu->startcodeLen], nalu->len - nalu->startcodeLen);
        nalu->nalType = m_hH264->nal->nal_unit_type;
        nalu->sliceType = m_hH264->sh->slice_type;
        m_hH264->sh->read_slice_type = 0;
    }else{
        m_hH265->sh->read_slice_type = 1;
        h265_read_nal_unit(m_hH265, &buffer[nalu->startcodeLen], nalu->len - nalu->startcodeLen);
        nalu->nalType = m_hH265->nal->nal_unit_type;
        nalu->sliceType = m_hH265->sh->slice_type;
        m_hH265->sh->read_slice_type = 0;
    }
    fseek(fp, pos, SEEK_SET);
    free(buffer);
    buffer = nullptr;
    return nalu->len;
}

int FileParse::findFirstNALU(FILE* fp, int* startcodeLenght)
{
    int found = 0;
    int info2 = 0;
    int info3 = 0;
    int pos = 0;
    int startcode_len = 0;
    unsigned char *buffer = nullptr;

    if ((buffer = (unsigned char*)calloc(MAX_NAL_SIZE, sizeof(char))) == nullptr)
        printf ("Could not allocate buffer memory\n");

    while (!found && !feof(fp))
    {
        buffer[pos++] = fgetc(fp);//读一个字节到BUF中

        info3 = findStartcode4(&buffer[pos-4]);//判断是否为0x00000001
        if(info3 != 1)
        {
            info2 = findStartcode3(&buffer[pos-3]);//判断是否为0x000001
            if (info2)
            {
                startcode_len = 3;
            }
        }
        else
        {
            startcode_len = 4;
        }

        found = (info2 == 1 || info3 == 1);
        if (pos >= MAX_NAL_SIZE)
        {
            free(buffer);
            buffer = nullptr;
            return -1;
        }
    }

    // 文件指针要恢复
    fseek(fp, -startcode_len, SEEK_CUR);

    free(buffer);
    buffer = nullptr;
    if (startcodeLenght != nullptr)
        *startcodeLenght = startcode_len;

    return pos - startcode_len;
}

FileParse::FileType FileParse::judeVideoFile(const char* filename)
{
    FileType type = FILE_H264; // default
    if(GetFileType(filename) != FILE_UNK){
        type = GetFileType(filename);
    }else{
        // read content 
        FILE* fp = nullptr;
        int offset = 0;
        int startcode = 0;
        unsigned char nalHader = 0;
        unsigned char nalType = 0;

        fp = fopen(filename, "r+b");
        offset = findFirstNALU(fp, &startcode);
        if (offset < 0)
        {
            return FILE_UNK;
        }
        fseek(fp, offset+startcode, SEEK_SET);
        int ret = fread((void*)&nalHader,1,1,fp);
        if(ret != 1){
            parse_error("fread Error");
            return FILE_UNK;
        }
        // check h264 first...
        nalType = nalHader & 0x1f; // 5 bit
        if (nalType > 0 && nalType < 22) // ok
        {
            type = FILE_H264;
        }
        else
        {
            // not h264, then check h265...
            nalType = (nalHader>>1) & 0x3f; // 6 bit
            if (nalType >= 0 && nalType <= 47) // ok
            {
                type = FILE_H265;
            }
        }
        fclose(fp);
        fp = nullptr;
    }

    return type;
}
